package com.bungarus.cluster;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;
import org.apache.curator.utils.ZKPaths;
import org.apache.log4j.Logger;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Created by tang on 2018/11/14.
 */
public class ClusterManager {
    Logger logger = Logger.getLogger(ClusterManager.class);

    private static final String BASE_URL = "/com/bungarus/";
    private static final String BASE_CLUSTERS_TAG = "/clusters";
    private static final String BASE_CLUSTER_TAG = "/cluster";
    private static final String BASE_CUR_CLUSTER_TAG = "/cur_cluster";
    private static final String BASE_NODE_PREFIX = "node-";

    /**
     * zk cluster connection
     */
    private static CuratorFramework clusterConnection;
    /**
     * whether zk ping scheduler started or not
     */
    private boolean started;
    /**
     * cluster tag
     */
    private String tag;

    public ClusterManager(String tag) {
        this(tag, ClusterConnection.DEFAULT_CONNECTION_STRING, ClusterConnection.DEFAULT_ZK_SESSION_TIMEOUT);
    }

    public ClusterManager(String tag, String connectStr) {
        this(tag, connectStr, ClusterConnection.DEFAULT_ZK_SESSION_TIMEOUT);
    }

    public ClusterManager(String tag, String connectStr, int sessionTimeout) {
        this.tag = tag;
        try {
            initCluster(connectStr, sessionTimeout);
        }
        catch (Exception e) {
            //throw new IllegalStateException(e);
            e.printStackTrace();
        }
    }

    private synchronized void initCluster(String connectStr, int sessionTimeout) throws Exception {
        if(null != clusterConnection){
            return;
        }
        clusterConnection = CuratorFrameworkFactory.builder()
                .connectString(connectStr)
                .connectionTimeoutMs(ClusterConnection.DEFAULT_ZK_CONNECTION_TIMEOUT)
                .sessionTimeoutMs(sessionTimeout)
                .retryPolicy(new ExponentialBackoffRetry(1000, 3))
                .build();
        clusterConnection.start();

        String clustersPath = ZKPaths.makePath(BASE_URL, this.tag, BASE_CLUSTERS_TAG);
        String clusterPath = ZKPaths.makePath(BASE_URL, this.tag, BASE_CLUSTER_TAG);
        String currentCluster = ZKPaths.makePath(BASE_URL, this.tag, BASE_CUR_CLUSTER_TAG);

        initCluster0(clustersPath, clusterPath, currentCluster);
    }

    private void initCluster0(String clustersPath, String clusterPath, String currentCluster) throws Exception {
        String idPath;
        Stat exists = clusterConnection.checkExists().forPath(currentCluster);
        if(null == exists) {
            clusterConnection.create()
                    .creatingParentContainersIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(currentCluster);

            exists = clusterConnection.checkExists().forPath(clustersPath);
            if(null == exists) {
                clusterConnection.create()
                        .creatingParentContainersIfNeeded()
                        .withMode(CreateMode.PERSISTENT)
                        .forPath(clustersPath);
            }
            idPath = clusterConnection.create()
                    //.withProtection()
                    .withMode(CreateMode.PERSISTENT_SEQUENTIAL)
                    .forPath(ZKPaths.makePath(clustersPath, "_"));
            idPath = ZKPaths.getNodeFromPath(idPath);

            exists = clusterConnection.checkExists().forPath(clusterPath);
            if(null == exists) {
                clusterConnection.create()
                        .creatingParentContainersIfNeeded()
                        .withMode(CreateMode.PERSISTENT)
                        .forPath(clusterPath);
            }
            if(idPath.length() > 0){
                clusterConnection.create()
                        .creatingParentContainersIfNeeded()
                        .withMode(CreateMode.PERSISTENT)
                        .forPath(ZKPaths.makePath(clusterPath, idPath));

                clusterConnection.setData().forPath(currentCluster, idPath.getBytes("UTF-8"));
            }
        }
    }

    /**
     * create a new cluster
     * @return the created cluster id prefixed with '_'
     * @throws Exception
     */
    public String createCluster() throws Exception {
        String clustersPath = ZKPaths.makePath(BASE_URL, this.tag, BASE_CLUSTERS_TAG, "_");
        String clusterPath = ZKPaths.makePath(BASE_URL, this.tag, BASE_CLUSTER_TAG);

        String idPath = clusterConnection.create()
                //.withProtection()
                .withMode(CreateMode.PERSISTENT_SEQUENTIAL)
                .forPath(clustersPath);
        idPath = ZKPaths.getNodeFromPath(idPath);

        clusterConnection.create()
                .withMode(CreateMode.PERSISTENT)
                .forPath(ZKPaths.makePath(clusterPath, idPath));

        return idPath;
    }

    public void setCurrentCluster(int clusterID) throws Exception{
        String currentCluster = ZKPaths.makePath(BASE_URL, this.tag, BASE_CUR_CLUSTER_TAG);
        String strClusterID = String.format("_%010d", clusterID);
        clusterConnection.setData().forPath(currentCluster, strClusterID.getBytes("UTF-8"));
    }

    public void setNodeData(NodeInfo node) throws Exception{
        String nodePath = ZKPaths.makePath(BASE_URL, node.getClusterBelongsTo(),
                BASE_CLUSTER_TAG, String.format("_%010d", node.getClusterIDBelongsTo()), BASE_NODE_PREFIX + String.format("%010d", node.getNodeId()));
        StringBuilder sb = new StringBuilder();
        sb.append(String.format("%010d", node.getNodeId()));
        sb.append(":");
        sb.append(node.getIp());
        sb.append(":");
        sb.append(node.getProtocol());
        sb.append(":");
        sb.append(node.getStatus());

        if(null != node.getMessagePorts()){
            Map<String, Integer> tagPorts = node.getMessagePorts();
            Iterator<String> tags = tagPorts.keySet().iterator();
            StringBuilder sbb = new StringBuilder("(");
            while(tags.hasNext()){
                String tag = tags.next();
                sbb.append(tag);
                sbb.append("|");
                sbb.append(tagPorts.get(tag));
                sbb.append(",");
            }
            if(sbb.length() > 1) { // remove comma
                sbb.substring(0, sbb.length() - 1);
            }
            sbb.append(")");
            sb.append(":");
            sb.append(sbb.toString());
        }
        clusterConnection.setData().forPath(nodePath, sb.toString().getBytes("UTF-8"));
    }

    public NodeInfo registerNode(Node node) throws RegisterNodeException {
        //using zk client to register the specified node to be one member of a cluster
        String currentCluster = ZKPaths.makePath(BASE_URL, node.getTag(), BASE_CUR_CLUSTER_TAG);
        try {
            String currentClusterID = new String(clusterConnection.getData().forPath(currentCluster), "UTF-8");
            String nodePath = ZKPaths.makePath(BASE_URL, node.getTag(), BASE_CLUSTER_TAG, currentClusterID, BASE_NODE_PREFIX);
            String returnPath = clusterConnection.create()
                    //.withProtection()
                    .withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
                    .forPath(nodePath);
            returnPath = ZKPaths.getNodeFromPath(returnPath);
            if(returnPath.length() <= BASE_NODE_PREFIX.length()){
                throw new RegisterNodeException();
            }

            String nodeID = returnPath.substring(BASE_NODE_PREFIX.length());
            StringBuilder sb = new StringBuilder(nodeID);
            sb.append(":");
            sb.append(node.getHostAddress());
            sb.append(":");
            sb.append(node.availableProtocol());
            sb.append(":");
            sb.append(Node.STATUS_INITIALIZING);
            sb.append(":");
            sb.append(node.tagPortString4Connection());

            clusterConnection.setData().forPath(nodePath + nodeID, sb.toString().getBytes("UTF-8"));

            NodeInfo nodeInfo = new NodeInfo();
            nodeInfo.setNodeId(Integer.valueOf(nodeID));
            nodeInfo.setClusterIDBelongsTo(Integer.valueOf(currentClusterID.substring(1)));
            nodeInfo.setClusterBelongsTo(node.getTag());
            nodeInfo.setProtocol(node.availableProtocol());
            nodeInfo.setStatus(Node.STATUS_INITIALIZING);
            nodeInfo.setIp(node.getHostAddress());

            if(node.tagPortString4Connection().length() >= 2){ // 2 stands for including only ()
                String str = node.tagPortString4Connection();
                String tagPortString = str.substring(1, str.length() - 1); // tag1|port1,tag2|port2
                String[] tagPortStrs = StringUtils.tokenizeToStringArray(tagPortString, ",");
                if(null != tagPortStrs && tagPortStrs.length > 0){
                    Map<String, Integer> map = new HashMap<>(0);
                    for(String tagPort : tagPortStrs){ //tag1|port1
                        String[] kv = StringUtils.split(tagPort, "|");
                        if(null != kv && kv.length >= 2) {
                            map.put(kv[0], Integer.valueOf(kv[1]));
                        }
                    }
                    if(map.size() > 0) {
                        nodeInfo.setMessagePorts(map);
                    }
                }
            }
            logger.info("=================== registered nodeInfo, cluster: " + node.getTag() + " detail: " + nodeInfo.toString());
            return nodeInfo;
        }
        catch (Exception e) {
            throw new RegisterNodeException(e);
        }
    }

    public List<ClusterInfo> clusters() throws Exception {
        return clusters(this.tag);
    }

    public List<ClusterInfo> clusters(String clusterTag) throws Exception {
        String clustersPath = ZKPaths.makePath(BASE_URL, clusterTag, BASE_CLUSTERS_TAG);
        List<String>clusterChildren = clusterConnection.getChildren().forPath(clustersPath);
        return clusterChildren.stream()
                .map(clusterPath -> {
                    ClusterInfo clusterInfo = new ClusterInfo();
                    clusterInfo.setTag(clusterTag);
                    clusterInfo.setClusterID(Integer.valueOf(clusterPath.substring(1)));
                    return clusterInfo;
                })
                .sorted((info1, info2) -> {
                    if (info1.getClusterID() == info2.getClusterID()) {
                        return 0;
                    }
                    else {
                        return info1.getClusterID() > info2.getClusterID() ? 1 : -1;
                    }
                })
                .collect(Collectors.toList());
    }

    public int getCurrentClusterID() throws Exception {
        String currentCluster = ZKPaths.makePath(BASE_URL, this.tag, BASE_CUR_CLUSTER_TAG);
        String currentClusterID = new String(clusterConnection.getData().forPath(currentCluster), "UTF-8");
        return Integer.valueOf(currentClusterID.substring(1));
    }

    public List<NodeInfo> nodes(int clusterId) throws Exception {
        return nodes(this.tag, clusterId);
    }

    public List<NodeInfo> nodes(String tag, int clusterId) throws Exception {
        String clusterPath = ZKPaths.makePath(BASE_URL, tag, BASE_CLUSTER_TAG, String.format("_%010d", clusterId));
        List<String> nodePaths = clusterConnection.getChildren().forPath(clusterPath);
        return nodePaths.stream()
                .map(nodePath -> {
                    try {
                        //node data format: nodeId:ip:protocol:status:(tag1|port1,tag2|port2)
                        String nodeDataStr = new String(clusterConnection.getData().forPath(ZKPaths.makePath(clusterPath, nodePath)), "UTF-8");
                        String[] nodeData = StringUtils.tokenizeToStringArray(nodeDataStr, NodeInfo.SEPARATOR);

                        NodeInfo nodeInfo = new NodeInfo();
                        nodeInfo.setNodeId(Integer.valueOf(nodeData[0]));
                        nodeInfo.setClusterIDBelongsTo(clusterId);
                        nodeInfo.setClusterBelongsTo(tag);
                        nodeInfo.setIp(nodeData[1]);
                        nodeInfo.setProtocol(Integer.valueOf(nodeData[2]));
                        nodeInfo.setStatus(Integer.valueOf(nodeData[3]));

                        String tagPortString = nodeData[4];
                        if (tagPortString.length() >= 2) { // 2 stands for including only ()
                            tagPortString = tagPortString.substring(1, tagPortString.length() - 1);
                            String[] tagPortStrs = StringUtils.tokenizeToStringArray(tagPortString, ",");
                            if (null != tagPortStrs && tagPortStrs.length > 0) {
                                Map<String, Integer> map = new HashMap<>(0);
                                for (String tagPort : tagPortStrs) { //tag1|port1
                                    String[] kv = StringUtils.split(tagPort, "|");
                                    if (null != kv && kv.length >= 2) {
                                        map.put(kv[0], Integer.valueOf(kv[1]));
                                    }
                                }
                                if (map.size() > 0) {
                                    nodeInfo.setMessagePorts(map);
                                }
                            }
                        }
                        return nodeInfo;
                    }
                    catch (Exception e) {
                        e.printStackTrace();
                    }
                    return null;
                })
                .collect(Collectors.toList());
    }

    public void close() {
        CloseableUtils.closeQuietly(clusterConnection);
    }
}
